import random
import matplotlib.pyplot as plt
import numpy as np 
import math
import time

#适应度
def fit(L):
    #初始相互攻击次数为0
    x=0
    #设置两个皇后递归循环
    for i1,j1 in enumerate(L,1):
        for i2,j2 in enumerate(L,1):
            #对比不同行的两个皇后
            if i2>i1:
                #同列或相互在对角线上都会引发相互攻击，X次数加1
                if j2==j1 or (abs(j2-j1)==i2-i1):x+=1
    return x
#生成初始种群
#生成随机候选解
def generate_one(N):
    #生成从1到N+1 步长为1的序列
    x=np.arange(1,N+1)
    #随机打乱序列
    np.random.shuffle(x)
    #返回随机打乱序列，与随机打乱序列的fit值
    #实际就是将每个候选解与适应度函数计算结果同时输出
    return (x,fit(x))

#定义生成初始种群的基因个数
def generate(N,NP):
    #按照传入的基因个数循环生成随机候选解
    return [generate_one(N) for _ in range(NP)]
#测试种群初始化函数
#输出5个8皇后问题候选解，作为初始种群
generate(8,5)

#定义突变，随机抽取一位并替换为一个随机数
def mutate(G):
    #在原基因列表中随机抽取一位
    x=random.randint(0,len(G)-1)
    #根据随机抽取位置索引，将基因列表中该索引号对应的数值替换为1到列表长度间的随机数
    G[x]=random.randint(1,len(G))
    return G

#定义交叉，输入两个基因
def crossover(G1,G2):
    x=random.randint(0,len(G1)) #从第一个基因中抽取随机点
    #在随机点截断两个基因
    #从G1抽取随机点前部位数，从G2抽取随机点后部位数
    #组合为新的基因
    print("随机点",x)
    return np.concatenate((G1[:x],G2[x:]))

#测试交叉函数
print(crossover([1,2,3,4],[5,6,7,8]))
print(mutate([1,2,3,4]))

#丢弃最差候选解
def discard_unfit(P):  #传入候选解[x,fit(x)]
    #按适应度函数排序
    #sort默认为升序，即fit值从小到大
    P.sort(key=lambda x:x[1])
    #只保留最优的1/3 fit值的基因，其余丢弃
    return P[:len(P)//3]

#该函数是从保留最优的1/3 fit值的基因中，按概率选择最优的2个基因索引
def choose_rand(P):
    #当前种群大小，即最优 1/3 的数量
    N=len(P[0][0])
    #适应度函数的最大值，即所有皇后之间都进行一次攻击
    mf=N*(N-1)//2
    #计算适应度最大值和最优种群（1/3）中每个基因的适应度值的差
    #注意：如果差值越大，意味着基因越优秀，下面计算概率的时候概率值越大
    z=[mf-x[1] for x in P]
    #计算上一步所有差值的和
    tf=sum(z)
    #计算每个差值相对于差值总和的概率（所有概率的和是1）
    p=np.random.choice(len(P),2,False,p=w)
    return p[0],p[1]

#生成新种群
def nxgeneration(P):
    #继承种群大小
    gen_size=len(P)
    #执行丢弃策略
    P=discard_unfit(P)
    #为突变与交叉增补三个基因位置
    P.extend(generate(len(P[0][0]),3))
    new_gen=[]
    for _ in range(gen_size):
        #随机生成两个新基因
        p1,p2=choose_rand(P)
        #对上面两个新基因执行交叉，成为第三个增补基因
        n=crossover(P[p1][0],P[p2][0])
        #随机概率>突变概率，则先执行突变
        #否则直接使用交叉后的基因计算fit值
        if random.random()<mutation_prob:
            n=mutate(n)
        nf=fit(n)
        #获得新一代种群
        new_gen.append((n,nf))
    return new_gen
#主函数，执行遗传算法
def genetic(N,pop_size=100):
    #种群初始化
    P=generate(N,pop_size)
    #提取初始化种群最小的适应度函数值
    mf=min([x[1] for x in P])
    n=0
    #如果最小的适应度函数值大于0
    while mf>0:
        n+=1
        #生成新种群，计算新种群中最小的fit值是否依然大于0
        #如果是，则继续循环
        p=nxgeneration(P)
        mf=min([x[1] for x in P])
    #直到新种群中最小的fit值等于0，while循环结束
    #根据最小fit值的索引提取对应基因（求解结果）
    mi=np.argmin([x[1] for x in P])
    return P[mi]
print("对于四皇后问题，最优解是：",genetic(4))